#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Python代码判题执行器
使用RestrictedPython进行安全代码执行
"""

import sys
import json
import time
import traceback
import ast
import signal
import resource
from io import StringIO
from contextlib import redirect_stdout, redirect_stderr


class TimeoutException(Exception):
    """超时异常"""
    pass


def timeout_handler(signum, frame):
    """超时信号处理器"""
    raise TimeoutException("代码执行超时")


def extract_function_name(code):
    """
    从代码中自动提取函数名
    使用AST解析，查找第一个函数定义
    """
    try:
        tree = ast.parse(code)
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                return node.name
        return "solution"  # 默认函数名
    except Exception:
        return "solution"


def parse_input_value(input_str):
    """
    智能解析输入值
    支持旧格式（纯文本）和新格式（JSON对象）
    """
    if not input_str or input_str.strip() == "":
        return None
    
    input_str = input_str.strip()
    
    # 尝试解析为JSON对象（新格式）
    try:
        data = json.loads(input_str)
        if isinstance(data, dict):
            # 新格式：{"type": "array", "value": [1,2,3], "params": [...]}
            if "params" in data:
                return data["params"]
            elif "value" in data:
                return [data["value"]]
            else:
                return [data]
        else:
            # JSON数组或基本类型
            return [data]
    except (json.JSONDecodeError, ValueError):
        pass
    
    # 尝试解析为Python字面量
    try:
        value = ast.literal_eval(input_str)
        return [value]
    except (ValueError, SyntaxError):
        pass
    
    # 旧格式：纯文本字符串
    return [input_str]


def parse_expected_output(output_str):
    """
    解析期望输出
    """
    if not output_str or output_str.strip() == "":
        return None
    
    output_str = output_str.strip()
    
    # 尝试解析为JSON
    try:
        data = json.loads(output_str)
        if isinstance(data, dict) and "value" in data:
            return data["value"]
        return data
    except (json.JSONDecodeError, ValueError):
        pass
    
    # 尝试解析为Python字面量
    try:
        return ast.literal_eval(output_str)
    except (ValueError, SyntaxError):
        pass
    
    # 纯文本
    return output_str


def create_safe_globals():
    """
    创建安全的全局命名空间
    只包含必要的内置函数
    """
    safe_builtins = {
        'abs': abs,
        'all': all,
        'any': any,
        'bool': bool,
        'chr': chr,
        'dict': dict,
        'enumerate': enumerate,
        'filter': filter,
        'float': float,
        'int': int,
        'isinstance': isinstance,
        'len': len,
        'list': list,
        'map': map,
        'max': max,
        'min': min,
        'ord': ord,
        'pow': pow,
        'print': print,
        'range': range,
        'reversed': reversed,
        'round': round,
        'set': set,
        'sorted': sorted,
        'str': str,
        'sum': sum,
        'tuple': tuple,
        'zip': zip,
        'True': True,
        'False': False,
        'None': None,
    }
    
    return {
        '__builtins__': safe_builtins,
        '__name__': '__main__',
        '__doc__': None,
    }


def execute_code_safely(code, function_name, test_input, time_limit=5):
    """
    安全执行用户代码
    
    Args:
        code: 用户代码
        function_name: 函数名
        test_input: 测试输入（列表形式）
        time_limit: 时间限制（秒）
    
    Returns:
        (result, error, execution_time)
    """
    # 设置超时
    signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(time_limit)
    
    # 设置内存限制（256MB）
    try:
        resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, 256 * 1024 * 1024))
    except Exception:
        pass  # 某些环境可能不支持
    
    start_time = time.time()
    result = None
    error = None
    
    try:
        # 创建安全的全局命名空间
        safe_globals = create_safe_globals()
        
        # 捕获标准输出和错误输出
        stdout_capture = StringIO()
        stderr_capture = StringIO()
        
        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
            # 编译并执行代码
            compiled_code = compile(code, '<user_code>', 'exec')
            exec(compiled_code, safe_globals)
            
            # 获取用户定义的函数
            if function_name not in safe_globals:
                raise NameError(f"函数 '{function_name}' 未定义")
            
            user_function = safe_globals[function_name]
            
            # 执行函数
            if test_input is None or len(test_input) == 0:
                result = user_function()
            elif len(test_input) == 1:
                result = user_function(test_input[0])
            else:
                result = user_function(*test_input)
        
        # 取消超时
        signal.alarm(0)
        
    except TimeoutException as e:
        error = str(e)
        signal.alarm(0)
    except MemoryError:
        error = "内存超限"
        signal.alarm(0)
    except Exception as e:
        error = f"{type(e).__name__}: {str(e)}"
        signal.alarm(0)
    
    execution_time = time.time() - start_time
    
    return result, error, execution_time


def compare_output(actual, expected):
    """
    比较实际输出和期望输出
    """
    # 处理None值
    if actual is None and expected is None:
        return True
    if actual is None or expected is None:
        return False
    
    # 类型转换：将结果转为可比较的格式
    def normalize(value):
        if isinstance(value, (list, tuple)):
            return list(value)
        elif isinstance(value, set):
            return sorted(list(value))
        elif isinstance(value, dict):
            return value
        elif isinstance(value, float):
            return round(value, 6)  # 浮点数精度
        else:
            return value
    
    actual_normalized = normalize(actual)
    expected_normalized = normalize(expected)
    
    return actual_normalized == expected_normalized


def judge(code_file, test_cases_file, function_name):
    """
    主判题函数
    
    Args:
        code_file: 代码文件路径
        test_cases_file: 测试用例JSON文件路径
        function_name: 函数名（可选，会自动检测）
    
    Returns:
        JSON格式的判题结果
    """
    result = {
        "status": "SE",
        "score": 0,
        "total_score": 100,
        "passed_cases": 0,
        "total_cases": 0,
        "execution_time": 0.0,
        "memory_usage": 0.0,
        "error_message": "",
        "test_results": []
    }
    
    try:
        # 读取代码
        with open(code_file, 'r', encoding='utf-8') as f:
            code = f.read()
        
        if not code.strip():
            result["error_message"] = "代码为空"
            return result
        
        # 自动检测函数名
        if not function_name or function_name == "solution":
            detected_name = extract_function_name(code)
            if detected_name:
                function_name = detected_name
        
        # 读取测试用例
        with open(test_cases_file, 'r', encoding='utf-8') as f:
            test_cases = json.load(f)
        
        if not test_cases:
            result["error_message"] = "没有测试用例"
            return result
        
        result["total_cases"] = len(test_cases)
        total_time = 0.0
        
        # 逐个执行测试用例
        for idx, test_case in enumerate(test_cases):
            case_result = {
                "case_id": idx + 1,
                "input": test_case.get("input", ""),
                "expected": test_case.get("expected", ""),
                "actual": "",
                "status": "SE",
                "time": 0.0,
                "memory": 0.0,
                "error": ""
            }
            
            try:
                # 解析输入
                test_input = parse_input_value(test_case.get("input", ""))
                expected_output = parse_expected_output(test_case.get("expected", ""))
                
                # 执行代码
                actual_output, error, exec_time = execute_code_safely(
                    code, function_name, test_input, time_limit=5
                )
                
                case_result["time"] = round(exec_time, 3)
                total_time += exec_time
                
                if error:
                    # 执行出错
                    if "超时" in error:
                        case_result["status"] = "TLE"
                    elif "内存" in error:
                        case_result["status"] = "MLE"
                    else:
                        case_result["status"] = "RE"
                    case_result["error"] = error
                else:
                    # 比较输出
                    case_result["actual"] = json.dumps(actual_output, ensure_ascii=False)
                    
                    if compare_output(actual_output, expected_output):
                        case_result["status"] = "AC"
                        result["passed_cases"] += 1
                    else:
                        case_result["status"] = "WA"
            
            except Exception as e:
                case_result["status"] = "SE"
                case_result["error"] = f"测试用例执行异常: {str(e)}"
            
            result["test_results"].append(case_result)
        
        # 计算总分
        if result["total_cases"] > 0:
            result["score"] = int((result["passed_cases"] / result["total_cases"]) * 100)
        
        result["execution_time"] = round(total_time, 3)
        
        # 确定最终状态
        if result["passed_cases"] == result["total_cases"]:
            result["status"] = "AC"
        elif result["passed_cases"] > 0:
            result["status"] = "PA"
        else:
            # 检查是否所有用例都是同一种错误
            statuses = [tc["status"] for tc in result["test_results"]]
            if all(s == "WA" for s in statuses):
                result["status"] = "WA"
            elif any(s == "TLE" for s in statuses):
                result["status"] = "TLE"
            elif any(s == "MLE" for s in statuses):
                result["status"] = "MLE"
            elif any(s == "RE" for s in statuses):
                result["status"] = "RE"
            else:
                result["status"] = "SE"
    
    except FileNotFoundError as e:
        result["error_message"] = f"文件未找到: {str(e)}"
    except json.JSONDecodeError as e:
        result["error_message"] = f"JSON解析错误: {str(e)}"
    except Exception as e:
        result["error_message"] = f"判题系统错误: {str(e)}\n{traceback.format_exc()}"
    
    return result


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print(json.dumps({
            "status": "SE",
            "error_message": "参数不足: 需要 code_file test_cases_file [function_name]"
        }))
        sys.exit(1)
    
    code_file = sys.argv[1]
    test_cases_file = sys.argv[2]
    function_name = sys.argv[3] if len(sys.argv) > 3 else "solution"
    
    result = judge(code_file, test_cases_file, function_name)
    print(json.dumps(result, ensure_ascii=False))
